﻿using System;
using System.Runtime.InteropServices;
using System.Text;
using UnityEngine;

namespace WhisperWrapper
{
    using whisper_token = Int32;

    public class Whisper
    {
#if UNITY_STANDALONE_WIN
        const string libraryName = "libwhisper.dll";
#elif UNITY_STANDALONE_LINUX
        const string libraryName = "libwhisper.so";
#elif UNITY_STANDALONE_OSX
        const string libraryName = "libwhisper.dylib";
#endif
        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern int whisper_decode(
            IntPtr ctx,
            IntPtr tokens,
            int n_tokens,
            int n_past,
            int n_threads
        );

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern int whisper_encode(
            IntPtr ctx,
            int offset,
            int n_threads
        );

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern void whisper_free(IntPtr ctx);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern void whisper_free_params(IntPtr parameters);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern void whisper_free_state(IntPtr state);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern int whisper_ctx_init_openvino_encoder(
            IntPtr ctx,
            string model_path,
            string device,
            string cache_dir
        );

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern IntPtr whisper_full_default_params_by_ref(WhisperSamplingStrategy strategy);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern WhisperFullParams whisper_full_default_params(WhisperSamplingStrategy strategy);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern IntPtr whisper_print_system_info();

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern int whisper_full(
            IntPtr ctx,
            WhisperFullParams parameters,
            IntPtr samples,
            int n_samples
        );

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern int whisper_full_with_state(
            IntPtr ctx,
            IntPtr state,
            WhisperFullParams parameters,
            IntPtr samples,
            int n_samples
        );


        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern int whisper_full_parallel(
            IntPtr ctx,
            WhisperFullParams parameters,
            IntPtr samples,
            int n_samples,
            int n_processors
        );

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern int whisper_full_n_segments(IntPtr ctx);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern int whisper_full_n_segments_from_state(IntPtr state);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern IntPtr whisper_full_get_segment_text(IntPtr ctx, int i_segment);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern IntPtr whisper_full_get_segment_text_from_state(IntPtr state, int i_segment);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern IntPtr whisper_init_from_file([MarshalAs(UnmanagedType.LPStr)] string path_model);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern IntPtr whisper_init_state(IntPtr context);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern IntPtr whisper_token_to_str(IntPtr ctx, whisper_token token);

        [DllImport(libraryName, CallingConvention = CallingConvention.Cdecl, CharSet = CharSet.Ansi)]
        static extern int whisper_tokenize(
            IntPtr ctx,
            string text,
            IntPtr tokens,
            int n_max_tokens
        );

        public static IntPtr ctx;
        public static WhisperFullParams wparams;

        public static void Initialize(string model, string language = "auto", string device = "CPU")
        {
            try
            {
                ctx = whisper_init_from_file(model);
                if (ctx == null)
                {
                    throw new Exception(string.Format("Error: failed to initialize whisper context {0}?\n", model));
                }

                whisper_ctx_init_openvino_encoder(ctx, null, device, null);
                var pref = whisper_full_default_params_by_ref(WhisperSamplingStrategy.WHISPER_SAMPLING_GREEDY);
                wparams = Marshal.PtrToStructure<WhisperFullParams>(pref);
                whisper_free_params(pref);
                IntPtr? lang = Marshal.StringToHGlobalAnsi(language);
                wparams.language = lang.Value;
                var onNewSegmentDelegate = new WhisperNewSegmentCallback(WhisperEventHandler.OnNewSegment);
                wparams.new_segment_callback = Marshal.GetFunctionPointerForDelegate(onNewSegmentDelegate);
            }
            catch (Exception e)
            {
                throw new Exception(e.Message + "\n" + e.StackTrace);
            }
        }

        public static string OnNewSegment(IntPtr ctx, IntPtr state, int n_new, IntPtr user_data)
        {
            int segments = whisper_full_n_segments_from_state(state);
            //Debug.Log(segments);
            string text = "";
            for (int i = 0; i < segments; ++i)
            {
                IntPtr ptxt = whisper_full_get_segment_text_from_state(state, i);
                text += StringFromNativeUtf8(ptxt);
                //Debug.Log(text);
            }
            return text;
        }

        public unsafe static void Inference(float[] data)
        {
            try
            {
                fixed (float* ptr = data)
                {
                    whisper_full_parallel(ctx, wparams, (IntPtr)ptr, data.Length, 1);
                    //var state = whisper_init_state(ctx);
                    //try
                    //{
                    //    whisper_full_with_state(ctx, state, wparams, (IntPtr)ptr, data.Length);
                    //}
                    //finally
                    //{
                    //    whisper_free_state(state);
                    //}
                }
            }
            catch (Exception e)
            {
                throw new Exception(e.Message + "\n" + e.StackTrace);
            }
        }

        public static void Destroy()
        {
            whisper_free(ctx);
        }

        public static string StringFromNativeUtf8(IntPtr nativeUtf8)
        {
            // check input null
            if (nativeUtf8 == IntPtr.Zero)
                return null;

            // find null terminator
            var len = 0;
            while (Marshal.ReadByte(nativeUtf8, len) != 0) ++len;

            // check empty string
            if (len == 0)
                return "";

            // copy buffer from beginning to null position 
            var buffer = new byte[len];
            Marshal.Copy(nativeUtf8, buffer, 0, buffer.Length);
            return Encoding.UTF8.GetString(buffer);
        }
    }
}
